#!/usr/bin/env python3
# coding: utf-8
# File: transfer_data.py

from src.util import get_project_root
import os

PROJECT_ROOT = get_project_root()


class TransferData:
    def __init__(self, save_filename):
        self.label_dict = {
            "治疗": "TREATMENT",
            "身体部位": "BODY",
            "症状和体征": "SIGNS",
            "疾病和诊断": "DISEASE",
            "检查和检验": "CHECK",
            "推断项": "INFER",
            "分割子句": "SECTION",
            "结果": "RESULT"}

        self.cate_dict = {
            'O': 0,
            'TREATMENT-I': 1,
            'TREATMENT-B': 2,
            'BODY-B': 3,
            'BODY-I': 4,
            'SIGNS-I': 5,
            'SIGNS-B': 6,
            'CHECK-B': 7,
            'CHECK-I': 8,
            'DISEASE-I': 9,
            'DISEASE-B': 10,
            'INFER-B': 11,
            'INFER-I': 12,
            'SECTION-B': 13,
            'SECTION-I': 14,
            'RESULT-B': 15,
            'RESULT-I': 16,
            'START': 17,
            'STOP': 18
        }

        self.origin_path = PROJECT_ROOT + '//data//'
        self.train_filepath = '{}//data//{}'.format(PROJECT_ROOT, save_filename)

        return

    def transfer(self, dir_set):
        '''
        将生成的人工标注文件，生成用于训练LSTM的训练标注文件train.txt
        :param dir_set:
        :return:
        '''
        f = open(self.train_filepath, 'w', encoding="UTF-8")
        count = 0
        for dir_name in dir_set:
            for root, dirs, files in os.walk(self.origin_path + dir_name):
                for file in files:
                    filepath = os.path.join(root, file)
                    if 'original' not in filepath:
                        continue
                    label_filepath = filepath.replace('.txtoriginal', '')
                    print(filepath, '\t\t', label_filepath)
                    content = open(filepath, encoding="UTF-8").read().strip()
                    res_dict = {}
                    for line in open(label_filepath, encoding="UTF-8"):
                        res = line.strip().split('	')
                        start = int(res[1])
                        end = int(res[2])
                        label = res[3]
                        label_id = self.label_dict.get(label)
                        for i in range(start, end + 1):
                            if i == start:
                                label_cate = label_id + '-B'
                            else:
                                label_cate = label_id + '-I'
                            res_dict[i] = label_cate

                    for indx, char in enumerate(content):
                        char_label = res_dict.get(indx, 'O')
                        print(char, char_label)
                        f.write(char + '\t' + char_label + '\n')

                    count += 1
        print("Total number of files processed: ", count)
        f.close()
        return
